from langgraph.graph import StateGraph, START, END
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from typing_extensions import TypedDict
from item_test import lab_item_test, lab_origin_message
from register import register, contract_access
from report import (
    report_archive,
    report_generate,
    report_print,
    report_review,
    report_send,
)
from sampling import sampling_item, sampling_plan

from sample import sample_receive, sample_out, sample_deal
from IPython.display import display, Image


class State(TypedDict):
    ###业务主键
    busId: int
    ###抽样主键
    samplingId: int
    ###样品主键
    sampleId: int
    ###检测项目主键
    itemTestId: int
    ###流程状态
    bus_node: str
    ###流程状态
    next_bus_node: str


class Business:
    def __init__(self):

        _builder = StateGraph(State)
        _builder_with_node = self._init_node(_builder)
        _graph = self._init_register_edge(_builder_with_node)
        self._graph = _graph

    def _init_node(self, builder):

        ##抽样部分下单
        ####################################################
        builder.add_node("samplingPlan", sampling_plan.SamplingPlan)
        builder.add_node("samplingItem", sampling_item.SamplingItem)

        ##任务下单部分
        ####################################################
        builder.add_node("register", register.Register)
        builder.add_node("contractAccess", contract_access.ContractAccess)

        ##样品部分
        ####################################################
        builder.add_node("sampleReceive", sample_receive.SampleReceive)
        builder.add_node("sampleOut", sample_out.SampleOut)
        builder.add_node("sampleDeal", sample_deal.SampleDeal)

        # ##检测项目部分
        # ####################################################
        builder.add_node("labItemTest", lab_item_test.LabItemTest)
        builder.add_node("labOriginMessage", lab_origin_message.LabOriginMessage)

        # ##报告部分
        # ####################################################
        builder.add_node("reportGenerate", report_generate.ReportGenerate)
        builder.add_node("reportReview", report_review.ReportReview)
        builder.add_node("reportPrint", report_print.ReportPrint)
        builder.add_node("reportSend", report_send.ReportSend)
        builder.add_node("reportArchive", report_archive.ReportArchive)
        return builder

    def _register_sample_edge(state):
        if "acc" == state["next_bus_node"]:
            return "contractAccess"
        else:
            return "sampleReceive"

    def _lab_item_edge(state):
        if "report" == state["next_bus_node"]:
            return "reportGenerate"
        else:
            return "labOriginMessage"

    def _report_back_item(state):
        if "lab" == state["next_bus_node"]:
            return "labItemTest"
        elif "origin" == state["next_bus_node"]:
            return "labOriginMessage"
        else:
            return "reportReview"

    def _report_review(state):
        if "print" == state["next_bus_node"]:
            return "reportPrint"
        else:
            return "reportGenerate"

    def _report_print(state):
        if "send" == state["next_bus_node"]:
            return "reportSend"
        else:
            return "reportGenerate"

    def _report_send(state):
        if "archive" == state["next_bus_node"]:
            return "reportArchive"
        else:
            return "reportGenerate"

    def _report_archive(state):
        if "report" == state["next_bus_node"]:
            return "reportGenerate"
        else:
            return END

    def _init_register_edge(self, builder):
        builder.set_entry_point("samplingPlan")
        builder.add_edge("samplingPlan", "samplingItem")
        builder.add_edge("samplingItem", "register")
        builder.add_edge(START, "register")
        builder.add_conditional_edges(
            "register", self._register_sample_edge, ["contractAccess", "sampleReceive"]
        )
        builder.add_edge("contractAccess", "register")
        builder.add_edge("contractAccess", "sampleReceive")
        builder.add_edge("sampleReceive", "register")
        builder.add_edge("sampleReceive", "sampleOut")
        builder.add_edge("sampleOut", "sampleDeal")
        builder.add_edge("sampleDeal", END)
        builder.add_edge("sampleReceive", "labItemTest")
        builder.add_conditional_edges(
            "labItemTest", self._lab_item_edge, ["reportGenerate", "labOriginMessage"]
        )
        builder.add_edge("labOriginMessage", "reportGenerate")
        builder.add_conditional_edges(
            "reportGenerate",
            self._report_back_item,
            ["labItemTest", "labOriginMessage", "reportReview"],
        )
        builder.add_conditional_edges(
            "reportReview", self._report_review, ["reportGenerate", "reportPrint"]
        )
        builder.add_conditional_edges(
            "reportPrint", self._report_print, ["reportGenerate", "reportSend"]
        )
        builder.add_conditional_edges(
            "reportSend", self._report_send, ["reportGenerate", "reportArchive"]
        )
        builder.add_conditional_edges(
            "reportArchive", self._report_archive, ["reportArchive", END]
        )
        _graph = builder.compile()
        return _graph

    def __call__(self):
        return display(Image(self._graph.get_graph().draw_mermaid_png()))


if __name__ == "__main__":
    business = Business()
    business()
